#########################################################################
## https://greenleaflab.github.io/chromVAR/articles/Articles/Applications.html#differential-accessibility-and-variability

library(Seurat)
library(chromVAR)
library(SummarizedExperiment)
library(parallel)
library(dplyr)
library(data.table)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--peak_gene_file"), type = "character"),
    make_option(c("--rnaData_file"), type = "character"),
    make_option(c("--atacData_file"), type = "character"),
    make_option(c("--cluster"), type = "character"),
    make_option(c("--data_base"), type = "character"),
    make_option(c("--cor_gene_peak"), type = "character"),
    make_option(c("--cell_name`"), type = "character"),
    make_option(c("--cpu"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){

    peak_gene_file <- "~/20231121_singleMuti/results/cluster_all_result/cluster2/correlation/cluster2_correlation_magic.Rdata"
    rnaData_file <- "~/20231121_singleMuti/input/testis_combined.Rdata"
    atacData_file <- "~/20231121_singleMuti/results/motif/cluster2/jaspar/cluster2.chromvar.rda"
    cluster <- "cluster2"

    cell_name <- "Sperm"
    out_path <- "~/20231121_singleMuti/results/motif/cluster2/cisbp"
    data_base <- "jaspar"
    cpu <- 10
    cor_gene_peak <- 0.5
    
}


###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

peak_gene_file <- opt$peak_gene_file
rnaData_file <- opt$rnaData_file
atacData_file <- opt$atacData_file
cell_name <- opt$cell_name
cluster <- opt$cluster
data_base <- opt$data_base
cpu <- opt$cpu
cor_gene_peak <- opt$cor_gene_peak
out_path <- opt$out_path

dir.create( paste0(out_path , "/TF-Gene_LinkageSore-Correlation") , recursive = T )

##########################################################################################

atac_list <- load(atacData_file)
rna_list <- load(rnaData_file)
peak_gene_list <- load(peak_gene_file)
cluster <- gsub( "cluster" , "" , cluster )

##########################################################################################
#### 1、计算TF的活性和基因表达的相关性
#### 2、根据peak和基因的关系以及peak和TF的关系，计算TF在该基因的激活程度占TF总激活程度的比例
## for all linked peaks, sum of motif score scaled by peak-to-gene link correlation


## 提取对应cluster的表达
#testis_combined$seurat_clusters <- ifelse( testis_combined$seurat_clusters %in% c(12,13) , 
#    17 ,  testis_combined$seurat_clusters)
if(cluster == 17){
   use_rna <- subset( testis_combined , seurat_clusters %in% c(12,13) )
}else{
   use_rna <- subset( testis_combined , seurat_clusters == cluster )
}

## 提取motif在细胞的z-score
tf_score <- deviationScores(dev)

## 提取peak和其对应的靶基因
peak_targetgene <- subset( cluster_cor , correlation > cor_gene_peak )

## 提取存在peak调控的基因
gene_list <- unique(peak_targetgene$V11)
   
## chromvar会对peak的起始位置+1，这里剪掉
peak_dat_names <- data.frame(motif_ix_scores@rowRanges)
peak_dat_names <- paste0( peak_dat_names$seqnames , ":" , peak_dat_names$start-1 , "-" , peak_dat_names$end )

## 计算TF的活性评分和基因表达的相关性
for( gene in gene_list ){
    print(gene)

    ## 提取该基因在所有细胞的表达情况
    tmp_exp <- use_rna@assays$MAGIC_RNA@data[gene,]
	names(tmp_exp) <- sapply( strsplit(names(tmp_exp) , "_"  ) , "[" , 2 )
    
	## 提取共有的细胞
	use_cell <- names(tmp_exp)[names(tmp_exp) %in% colnames(tf_score)]
	tmp_exp <- tmp_exp[use_cell]
	tf_score_use <- tf_score[,use_cell]

    ## 提取关注基因与哪些peak相连
    tmp_peak_targetgene <- subset( peak_targetgene , V11 == gene )
    tmp_peak_targetgene$peak_id <- paste0( tmp_peak_targetgene$V1 , ":" , tmp_peak_targetgene$V2 , "-" , tmp_peak_targetgene$V3 )

	result_tmp <- c()

    ## 按照每个转录因子计算
    result_tmp <- Reduce(function(x,y)bind_rows( x , y ),mclapply(1:nrow(tf_score_use) , function(i){

        tf_name <- rownames(tf_score_use)[i]
        print(tf_name)

        #### TF的活性程度和基因表达的相关性
		correlation <- cor.test(tf_score_use[i, ], tmp_exp , method = "pearson")

        #### 计算linkage_sore
        tf_index <- which( colnames(motif_ix_scores) == tf_name )
        tmp_tf_score <- as.numeric(assay(motif_ix_scores[,tf_index]))
        names(tmp_tf_score) <- peak_dat_names
        tmp_tf_score <- tmp_tf_score[tmp_tf_score>0]

        ## 与基因表达相关的peak
        tmp_cor_peak <- tmp_tf_score[names(tmp_tf_score) %in% tmp_peak_targetgene$peak_id]
        tmp_cor_peak <- data.frame( peak_id = names(tmp_cor_peak) , tf_score = tmp_cor_peak )
        tmp_tf_score <- merge( tmp_peak_targetgene , tmp_cor_peak , by = "peak_id" )
        linkage_sore <- sum( tmp_tf_score$tf_score*(tmp_tf_score$correlation^2) )

        tmp_res <- data.frame( tf = tf_name , gene = gene , 
            median_zscore_tf = median(tf_score_use[i,]) , median_expression_gene = median(tmp_exp) ,
            cor_person = correlation$estimate , p_person = correlation$p.value , linkage_sore = linkage_sore )

	   tmp_res

    },mc.cores=cpu))

    ## 基因名里面存在"THRA1/BTR"，离谱
    out_name <- paste0(out_path , "/TF-Gene_LinkageSore-Correlation/" , gsub( "/" , "-" , gene) , ".tsv")
    write.table( result_tmp , out_name , sep = "\t" , row.names = F , quote = F )
}




